{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Transfer Learning\n",
    "\n",
    "Most of the time you won't want to train a whole convolutional network yourself. Modern ConvNets training on huge datasets like ImageNet take weeks on multiple GPUs. \n",
    "> Instead, most people use a pretrained network either as a fixed feature extractor, or as an initial network to fine tune. \n",
    "\n",
    "In this notebook, you'll be using [VGGNet](https://arxiv.org/pdf/1409.1556.pdf) trained on the [ImageNet dataset](http://www.image-net.org/) as a feature extractor. Below is a diagram of the VGGNet architecture, with a series of convolutional and maxpooling layers, then three fully-connected layers at the end that classify the 1000 classes found in the ImageNet database.\n",
    "\n",
    "<img src=\"notebook_ims/vgg_16_architecture.png\" width=700px>\n",
    "\n",
    "VGGNet is great because it's simple and has great performance, coming in second in the ImageNet competition. The idea here is that we keep all the convolutional layers, but **replace the final fully-connected layer** with our own classifier. This way we can use VGGNet as a _fixed feature extractor_ for our images then easily train a simple classifier on top of that. \n",
    "* Use all but the last fully-connected layer as a fixed feature extractor.\n",
    "* Define a new, final classification layer and apply it to a task of our choice!\n",
    "\n",
    "You can read more about transfer learning from [the CS231n Stanford course notes](http://cs231n.github.io/transfer-learning/).\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Flower power\n",
    "\n",
    "Here we'll be using VGGNet to classify images of flowers. We'll start, as usual, by importing our usual resources. And checking if we can train our model on GPU.\n",
    "\n",
    "### Download Data\n",
    "\n",
    "Download the flower data from [this link](https://s3.amazonaws.com/video.udacity-data.com/topher/2018/September/5baa60a0_flower-photos/flower-photos.zip), save it in the home directory of this notebook and extract the zip file to get the directory `flower_photos/`. **Make sure the directory has this exact name for accessing data: flower_photos**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "import torchvision\n",
    "from torchvision import datasets, models, transforms\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# check if CUDA is available\n",
    "train_on_gpu = torch.cuda.is_available()\n",
    "\n",
    "if not train_on_gpu:\n",
    "    print('CUDA is not available.  Training on CPU ...')\n",
    "else:\n",
    "    print('CUDA is available!  Training on GPU ...')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load and Transform our Data\n",
    "\n",
    "We'll be using PyTorch's [ImageFolder](https://pytorch.org/docs/stable/torchvision/datasets.html#imagefolder) class which makes it very easy to load data from a directory. For example, the training images are all stored in a directory path that looks like this:\n",
    "```\n",
    "root/class_1/xxx.png\n",
    "root/class_1/xxy.png\n",
    "root/class_1/xxz.png\n",
    "\n",
    "root/class_2/123.png\n",
    "root/class_2/nsdf3.png\n",
    "root/class_2/asd932_.png\n",
    "```\n",
    "\n",
    "Where, in this case, the root folder for training is `flower_photos/train/` and the classes are the names of flower types."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# define training and test data directories\n",
    "data_dir = 'flower_photos/'\n",
    "train_dir = os.path.join(data_dir, 'train/')\n",
    "test_dir = os.path.join(data_dir, 'test/')\n",
    "\n",
    "# classes are folders in each directory with these names\n",
    "classes = ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Transforming the Data\n",
    "\n",
    "When we perform transfer learning, we have to shape our input data into the shape that the pre-trained model expects. VGG16 expects `224`-dim square images as input and so, we resize each flower image to fit this mold."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# load and transform data using ImageFolder\n",
    "\n",
    "# VGG-16 Takes 224x224 images as input, so we resize all of them\n",
    "data_transform = transforms.Compose([transforms.RandomResizedCrop(224), \n",
    "                                      transforms.ToTensor()])\n",
    "\n",
    "train_data = datasets.ImageFolder(train_dir, transform=data_transform)\n",
    "test_data = datasets.ImageFolder(test_dir, transform=data_transform)\n",
    "\n",
    "# print out some data stats\n",
    "print('Num training images: ', len(train_data))\n",
    "print('Num test images: ', len(test_data))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### DataLoaders and Data Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# define dataloader parameters\n",
    "batch_size = 20\n",
    "num_workers=0\n",
    "\n",
    "# prepare data loaders\n",
    "train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, \n",
    "                                           num_workers=num_workers, shuffle=True)\n",
    "test_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size, \n",
    "                                          num_workers=num_workers, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Visualize some sample data\n",
    "\n",
    "# obtain one batch of training images\n",
    "dataiter = iter(train_loader)\n",
    "images, labels = dataiter.next()\n",
    "images = images.numpy() # convert images to numpy for display\n",
    "\n",
    "# plot the images in the batch, along with the corresponding labels\n",
    "fig = plt.figure(figsize=(25, 4))\n",
    "for idx in np.arange(20):\n",
    "    ax = fig.add_subplot(2, 20/2, idx+1, xticks=[], yticks=[])\n",
    "    plt.imshow(np.transpose(images[idx], (1, 2, 0)))\n",
    "    ax.set_title(classes[labels[idx]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Define the Model\n",
    "\n",
    "To define a model for training we'll follow these steps:\n",
    "1. Load in a pre-trained VGG16 model\n",
    "2. \"Freeze\" all the parameters, so the net acts as a fixed feature extractor \n",
    "3. Remove the last layer\n",
    "4. Replace the last layer with a linear classifier of our own\n",
    "\n",
    "**Freezing simply means that the parameters in the pre-trained model will *not* change during training.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Load the pretrained model from pytorch\n",
    "vgg16 = models.vgg16(pretrained=True)\n",
    "\n",
    "# print out the model structure\n",
    "print(vgg16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "print(vgg16.classifier[6].in_features) \n",
    "print(vgg16.classifier[6].out_features) \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Freeze training for all \"features\" layers\n",
    "for param in vgg16.features.parameters():\n",
    "    param.requires_grad = False\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "### Final Classifier Layer\n",
    "\n",
    "Once you have the pre-trained feature extractor, you just need to modify and/or add to the final, fully-connected classifier layers. In this case, we suggest that you replace the last layer in the vgg classifier group of layers. \n",
    "> This layer should see as input the number of features produced by the portion of the network that you are not changing, and produce an appropriate number of outputs for the flower classification task.\n",
    "\n",
    "You can access any layer in a pretrained network by name and (sometimes) number, i.e. `vgg16.classifier[6]` is the sixth layer in a group of layers named \"classifier\".\n",
    "\n",
    "#### TODO: Replace the last fully-connected layer with one that produces the appropriate number of class scores."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import torch.nn as nn## TODO: add a last linear layer  that maps n_inputs -> 5 flower classes\n",
    "## new layers automatically have requires_grad = True\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# after completing your model, if GPU is available, move the model to GPU\n",
    "if train_on_gpu:\n",
    "    vgg16.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "### Specify [Loss Function](http://pytorch.org/docs/stable/nn.html#loss-functions) and [Optimizer](http://pytorch.org/docs/stable/optim.html)\n",
    "\n",
    "Below we'll use cross-entropy loss and stochastic gradient descent with a small learning rate. Note that the optimizer accepts as input _only_ the trainable parameters `vgg.classifier.parameters()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import torch.optim as optim\n",
    "\n",
    "# specify loss function (categorical cross-entropy)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "# specify optimizer (stochastic gradient descent) and learning rate = 0.001\n",
    "optimizer = optim.SGD(vgg16.classifier.parameters(), lr=0.001)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Training\n",
    "\n",
    "Here, we'll train the network.\n",
    "\n",
    "> **Exercise:** So far we've been providing the training code for you. Here, I'm going to give you a bit more of a challenge and have you write the code to train the network. Of course, you'll be able to see my solution if you need help."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# number of epochs to train the model\n",
    "n_epochs = 2\n",
    "\n",
    "## TODO complete epoch and training batch loops\n",
    "## These loops should update the classifier-weights of this model\n",
    "## And track (and print out) the training loss over time\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Testing\n",
    "\n",
    "Below you see the test accuracy for each flower class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# track test loss \n",
    "# over 5 flower classes\n",
    "test_loss = 0.0\n",
    "class_correct = list(0. for i in range(5))\n",
    "class_total = list(0. for i in range(5))\n",
    "\n",
    "vgg16.eval() # eval mode\n",
    "\n",
    "# iterate over test data\n",
    "for data, target in test_loader:\n",
    "    # move tensors to GPU if CUDA is available\n",
    "    if train_on_gpu:\n",
    "        data, target = data.cuda(), target.cuda()\n",
    "    # forward pass: compute predicted outputs by passing inputs to the model\n",
    "    output = vgg16(data)\n",
    "    # calculate the batch loss\n",
    "    loss = criterion(output, target)\n",
    "    # update  test loss \n",
    "    test_loss += loss.item()*data.size(0)\n",
    "    # convert output probabilities to predicted class\n",
    "    _, pred = torch.max(output, 1)    \n",
    "    # compare predictions to true label\n",
    "    correct_tensor = pred.eq(target.data.view_as(pred))\n",
    "    correct = np.squeeze(correct_tensor.numpy()) if not train_on_gpu else np.squeeze(correct_tensor.cpu().numpy())\n",
    "    # calculate test accuracy for each object class\n",
    "    for i in range(batch_size):\n",
    "        label = target.data[i]\n",
    "        class_correct[label] += correct[i].item()\n",
    "        class_total[label] += 1\n",
    "\n",
    "# calculate avg test loss\n",
    "test_loss = test_loss/len(test_loader.dataset)\n",
    "print('Test Loss: {:.6f}\\n'.format(test_loss))\n",
    "\n",
    "for i in range(5):\n",
    "    if class_total[i] > 0:\n",
    "        print('Test Accuracy of %5s: %2d%% (%2d/%2d)' % (\n",
    "            classes[i], 100 * class_correct[i] / class_total[i],\n",
    "            np.sum(class_correct[i]), np.sum(class_total[i])))\n",
    "    else:\n",
    "        print('Test Accuracy of %5s: N/A (no training examples)' % (classes[i]))\n",
    "\n",
    "print('\\nTest Accuracy (Overall): %2d%% (%2d/%2d)' % (\n",
    "    100. * np.sum(class_correct) / np.sum(class_total),\n",
    "    np.sum(class_correct), np.sum(class_total)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "### Visualize Sample Test Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# obtain one batch of test images\n",
    "dataiter = iter(test_loader)\n",
    "images, labels = dataiter.next()\n",
    "images.numpy()\n",
    "\n",
    "# move model inputs to cuda, if GPU available\n",
    "if train_on_gpu:\n",
    "    images = images.cuda()\n",
    "\n",
    "# get sample outputs\n",
    "output = vgg16(images)\n",
    "# convert output probabilities to predicted class\n",
    "_, preds_tensor = torch.max(output, 1)\n",
    "preds = np.squeeze(preds_tensor.numpy()) if not train_on_gpu else np.squeeze(preds_tensor.cpu().numpy())\n",
    "\n",
    "# plot the images in the batch, along with predicted and true labels\n",
    "fig = plt.figure(figsize=(25, 4))\n",
    "for idx in np.arange(20):\n",
    "    ax = fig.add_subplot(2, 20/2, idx+1, xticks=[], yticks=[])\n",
    "    plt.imshow(np.transpose(images[idx], (1, 2, 0)))\n",
    "    ax.set_title(\"{} ({})\".format(classes[preds[idx]], classes[labels[idx]]),\n",
    "                 color=(\"green\" if preds[idx]==labels[idx].item() else \"red\"))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [default]",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  },
  "widgets": {
   "state": {},
   "version": "1.1.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
